#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Elasticsearch增强的Assistant
扩展qwen_agent的Assistant类，支持Elasticsearch检索
"""

import os
import logging
from typing import List, Dict, Any, Optional, Union
from qwen_agent.agents import Assistant
from qwen_agent.memory.elasticsearch_memory import ElasticsearchMemory

logger = logging.getLogger(__name__)


class ElasticsearchAssistant(Assistant):
    """
    支持Elasticsearch检索的Assistant
    """
    
    def __init__(self,
                 function_list: Optional[List[Union[str, Dict, Any]]] = None,
                 llm: Optional[Union[Dict, Any]] = None,
                 system_message: Optional[str] = None,
                 name: Optional[str] = None,
                 description: Optional[str] = None,
                 files: Optional[List[str]] = None,
                 use_elasticsearch: bool = True,
                 es_config: Optional[Dict[str, Any]] = None,
                 **kwargs):
        """
        初始化ElasticsearchAssistant
        
        Args:
            function_list: 工具列表
            llm: LLM配置
            system_message: 系统消息
            name: 助手名称
            description: 助手描述
            files: 文件列表
            use_elasticsearch: 是否使用Elasticsearch
            es_config: Elasticsearch配置
            **kwargs: 其他参数
        """
        self.use_elasticsearch = use_elasticsearch
        self.es_config = es_config or {}
        
        # 如果使用Elasticsearch，创建ES内存管理器
        if use_elasticsearch and files:
            logger.info("🔍 启用Elasticsearch检索模式")
            
            # 创建ES内存管理器
            self.es_memory = ElasticsearchMemory(
                llm=llm or {},
                files=files,
                es_config=self.es_config,
                **kwargs
            )
            
            # 获取ES检索工具
            es_tools = self.es_memory.get_tool_list()
            
            # 将ES工具添加到function_list
            if function_list is None:
                function_list = []
            
            # 添加ES检索工具（会覆盖默认的retrieval工具）
            function_list.extend(es_tools)
            
            # 不传递files给父类，避免创建默认的memory
            super().__init__(
                function_list=function_list,
                llm=llm,
                system_message=system_message,
                name=name,
                description=description,
                files=None,  # 不传递files
                **kwargs
            )
            
            # 手动设置memory为ES内存管理器
            self.memory = self.es_memory
            
        else:
            # 使用默认模式
            logger.info("📚 使用默认检索模式")
            super().__init__(
                function_list=function_list,
                llm=llm,
                system_message=system_message,
                name=name,
                description=description,
                files=files,
                **kwargs
            )
            self.es_memory = None
    
    def add_files(self, file_paths: List[str]) -> bool:
        """
        添加文件到索引
        
        Args:
            file_paths: 文件路径列表
            
        Returns:
            是否成功
        """
        if self.use_elasticsearch and self.es_memory:
            return self.es_memory.add_files(file_paths)
        else:
            # 使用默认方式添加文件
            if hasattr(self.memory, 'add_files'):
                return self.memory.add_files(file_paths)
            else:
                logger.warning("当前内存管理器不支持动态添加文件")
                return False
    
    def search_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """
        搜索文档
        
        Args:
            query: 查询字符串
            k: 返回结果数量
            
        Returns:
            搜索结果
        """
        if self.use_elasticsearch and self.es_memory:
            return self.es_memory.search_documents(query, k)
        else:
            # 使用默认检索
            if hasattr(self.memory, 'retrieve'):
                return self.memory.retrieve(query)
            else:
                logger.warning("当前内存管理器不支持文档搜索")
                return []
    
    def get_index_stats(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        if self.use_elasticsearch and self.es_memory:
            return self.es_memory.get_index_stats()
        else:
            return {"message": "未使用Elasticsearch"}
    
    def clear_index(self) -> bool:
        """清空索引"""
        if self.use_elasticsearch and self.es_memory:
            return self.es_memory.clear_index()
        else:
            logger.warning("未使用Elasticsearch，无法清空索引")
            return False
    
    def switch_to_elasticsearch(self, es_config: Optional[Dict[str, Any]] = None) -> bool:
        """
        切换到Elasticsearch模式
        
        Args:
            es_config: ES配置
            
        Returns:
            是否成功
        """
        try:
            if self.use_elasticsearch:
                logger.info("已经在使用Elasticsearch模式")
                return True
            
            # 获取当前文件列表
            current_files = getattr(self.memory, 'files', [])
            
            if not current_files:
                logger.warning("没有文件可以迁移到Elasticsearch")
                return False
            
            # 创建ES内存管理器
            self.es_config = es_config or self.es_config
            self.es_memory = ElasticsearchMemory(
                llm=self.llm.get_cfg() if hasattr(self.llm, 'get_cfg') else {},
                files=current_files,
                es_config=self.es_config
            )
            
            # 获取ES检索工具并添加到工具列表
            es_tools = self.es_memory.get_tool_list()
            
            # 移除现有的retrieval工具
            self.function_list = [tool for tool in self.function_list 
                                if not (hasattr(tool, 'name') and tool.name == 'retrieval')]
            
            # 添加ES检索工具
            self.function_list.extend(es_tools)
            
            # 切换内存管理器
            self.memory = self.es_memory
            self.use_elasticsearch = True
            
            logger.info("✅ 成功切换到Elasticsearch模式")
            return True
            
        except Exception as e:
            logger.error(f"❌ 切换到Elasticsearch模式失败: {str(e)}")
            return False
    
    def get_retrieval_mode(self) -> str:
        """获取当前检索模式"""
        return "Elasticsearch" if self.use_elasticsearch else "Default"


def create_elasticsearch_assistant(function_list: Optional[List[Union[str, Dict, Any]]] = None,
                                 llm: Optional[Union[Dict, Any]] = None,
                                 system_message: Optional[str] = None,
                                 name: Optional[str] = None,
                                 description: Optional[str] = None,
                                 files: Optional[List[str]] = None,
                                 es_config: Optional[Dict[str, Any]] = None,
                                 **kwargs) -> ElasticsearchAssistant:
    """
    创建支持Elasticsearch的Assistant
    
    Args:
        function_list: 工具列表
        llm: LLM配置
        system_message: 系统消息
        name: 助手名称
        description: 助手描述
        files: 文件列表
        es_config: ES配置
        **kwargs: 其他参数
        
    Returns:
        ElasticsearchAssistant实例
    """
    return ElasticsearchAssistant(
        function_list=function_list,
        llm=llm,
        system_message=system_message,
        name=name,
        description=description,
        files=files,
        use_elasticsearch=True,
        es_config=es_config,
        **kwargs
    )